
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math


class WeightedCrossEntropyLoss(nn.Module):
    """
    加权交叉熵损失函数，用于处理类别不平衡问题
    可以通过权重参数调整不同类别的重要性
    """
    def __init__(self, weight=None):
        super(WeightedCrossEntropyLoss, self).__init__()
        self.weight = torch.tensor(weight, dtype=torch.float) if weight is not None else None
    
    def forward(self, input, target):
        if self.weight is not None:
            # 确保权重在正确的设备上
            if self.weight.device != input.device:
                self.weight = self.weight.to(input.device)
        return F.cross_entropy(input, target, weight=self.weight)


def calc_class_weight(labels_count, mu=None):
    """
    基于训练集的类别分布计算类别权重，使用mu作为超参数
    
    参数:
    - labels_count: 每个类别的样本数列表
    - mu: 权重调整系数列表
    
    返回值:
    - class_weight: 类别权重列表
    """
    total = np.sum(labels_count)
    class_weight = dict()
    num_classes = len(labels_count)

    # 如果没有提供mu，则使用默认值
    if mu is None:
        factor = 1 / num_classes
        mu = [factor * 1.5, factor * 2, factor * 1.5, factor, factor * 1.5] # 这些配置仅适用于SLEEP-EDF-20数据集
    # 确保mu的长度与类别数匹配
    elif len(mu) != num_classes:
        print(f"警告: mu参数长度({len(mu)})与类别数({num_classes})不匹配，将截断或扩展mu参数")
        mu = mu[:num_classes] + [1.0] * (num_classes - len(mu))

    for key in range(num_classes):
        # 避免除以零错误
        if labels_count[key] == 0:
            print(f"警告: 类别{key}的样本数为0，将使用默认权重")
            class_weight[key] = 1.0  # 对于未出现的类别使用默认权重
        else:
            try:
                score = math.log(mu[key] * total / float(labels_count[key]))
                class_weight[key] = score if score > 1.0 else 1.0
                class_weight[key] = round(class_weight[key] * mu[key], 2)
            except Exception as e:
                print(f"计算类别{key}的权重时出错: {e}，将使用默认权重")
                class_weight[key] = 1.0

    class_weight = [class_weight[i] for i in range(num_classes)]
    return class_weight


from training.utils.stats_util import get_class_weights_from_stats

# 保留函数名以保持向后兼容性
def load_class_weights_from_stats(stats_file, normalize=True):
    """
    从统计文件加载类别权重（使用stats_util中的实现）
    
    参数:
    - stats_file: 统计文件路径
    - normalize: 是否归一化权重
    
    返回值:
    - class_weights: 计算得到的类别权重列表，如果加载失败则返回None
    """
    return get_class_weights_from_stats(stats_file)

def calculate_class_weights(dataset, mu=None, stats_file=None):
    """
    计算类别权重，仅从CSV统计文件加载并使用类别频率倒数计算权重
    
    参数:
    - dataset: 训练数据集对象（不再使用）
    - mu: 权重调整系数列表（不再使用）
    - stats_file: 统计文件路径，必须提供并有效
    
    返回值:
    - weights: 类别权重列表
    
    异常:
    - ValueError: 当统计文件不存在或无法从中获取权重时抛出
    """
    # 如果没有提供统计文件，直接报错
    if stats_file is None:
        raise ValueError("必须提供统计文件路径(stats_file)来获取类别权重")
    
    # 从统计文件加载权重
    class_weights = load_class_weights_from_stats(stats_file)
    if class_weights is not None:
        return class_weights
    else:
        # 当无法从统计文件获取权重时，直接报错
        raise ValueError(f"无法从统计文件'{stats_file}'中获取有效权重信息")


def create_loss_function(loss_type, loss_args=None, dataset=None, stats_file=None):
    """
    创建损失函数实例
    
    参数:
    - loss_type: 损失函数类型名称
    - loss_args: 损失函数的参数
    - dataset: 训练数据集对象（不再使用）
    - stats_file: 统计文件路径，必须提供并有效
    
    返回值:
    - criterion: 损失函数实例
    
    异常:
    - ValueError: 当损失函数不存在或无法加载时抛出
    - ValueError: 当需要权重但无法从统计文件获取时抛出
    """
    if loss_args is None:
        loss_args = {}
    
    # 检查是否需要自动计算类别权重
    if (loss_type == 'WeightedCrossEntropyLoss' and loss_args.get('weight') == 'auto') or \
       (loss_type == 'FocalLoss' and loss_args.get('alpha') == 'auto'):
        # 从loss_args中获取mu参数
        mu = loss_args.get('mu', None)
        # 必须提供统计文件
        if stats_file is None:
            raise ValueError("当weight或alpha设为'auto'时，必须提供统计文件路径(stats_file)")
        # 调用calculate_class_weights函数，此时该函数会严格从统计文件获取权重
        class_weights = calculate_class_weights(None, mu, stats_file)
        # 根据损失函数类型设置相应的参数
        if loss_type == 'WeightedCrossEntropyLoss':
            loss_args['weight'] = class_weights
        elif loss_type == 'FocalLoss':
            loss_args['alpha'] = class_weights
    
    try:
        # 检查自定义损失函数
        if loss_type == 'WeightedCrossEntropyLoss' or loss_type == 'FocalLoss':
            # 移除mu参数，因为这些损失函数不需要它
            loss_args_copy = loss_args.copy()
            if 'mu' in loss_args_copy:
                del loss_args_copy['mu']
            if loss_type == 'WeightedCrossEntropyLoss':
                criterion = WeightedCrossEntropyLoss(**loss_args_copy)
            else:  # FocalLoss
                criterion = FocalLoss(**loss_args_copy)
            print(f"成功加载自定义损失函数: {loss_type}")
            return criterion
        
        # 检查torch.nn中是否存在该损失函数
        elif hasattr(nn, loss_type):
            criterion_class = getattr(nn, loss_type)
            criterion = criterion_class(**loss_args)
            
            print(f"成功加载损失函数: {loss_type}")
            return criterion
        else:
            raise ValueError(f"PyTorch中不存在损失函数: {loss_type}")
    except Exception as e:
        raise ValueError(f"加载损失函数 {loss_type} 失败: {str(e)}")


class FocalLoss(nn.Module):
    """
    Focal Loss，用于处理类别不平衡问题
    减少容易分类的样本的权重，专注于难分类的样本
    
    参数:
    - alpha: 类别权重，用于类别不平衡
    - gamma: 聚焦参数，控制难易样本的权重调整程度
    """
    def __init__(self, alpha=None, gamma=2.0):
        super(FocalLoss, self).__init__()
        self.alpha = torch.tensor(alpha, dtype=torch.float) if alpha is not None else None
        self.gamma = gamma
    
    def forward(self, input, target):
        if self.alpha is not None:
            # 确保权重在正确的设备上
            if self.alpha.device != input.device:
                self.alpha = self.alpha.to(input.device)
        ce_loss = F.cross_entropy(input, target, weight=self.alpha, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()


class LabelSmoothingLoss(nn.Module):
    """
    标签平滑损失函数，用于防止模型过拟合
    
    参数:
    - classes: 类别数量
    - smoothing: 平滑系数，控制平滑程度
    - dim: 维度
    """
    def __init__(self, classes, smoothing=0.1, dim=-1):
        super(LabelSmoothingLoss, self).__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.cls = classes
        self.dim = dim
    
    def forward(self, input, target):
        input = F.log_softmax(input, dim=self.dim)
        with torch.no_grad():
            # 创建平滑标签
            true_dist = torch.zeros_like(input)
            true_dist.fill_(self.smoothing / (self.cls - 1))
            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        return torch.mean(torch.sum(-true_dist * input, dim=self.dim))